from PyQt5.QtWidgets import (QMainWindow, QSplitter, QAction, QFileDialog, 
                             QGraphicsView, QGraphicsScene, QLabel, QVBoxLayout, QWidget, QTabWidget, QMessageBox)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap
from core.alignment.mmaligment import AlignmentWidget
from core.annotator.annotator import AnnotationWidget
from core.image_processor.extract_frame import extract_frames
from gui.frame_extraction_panel import FrameExtractionPanel
from gui.image_view import ImageView
from gui.project_manager_panel import ProjectManagerPanel
import os
import glob

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("MMReID标注工具")
        self.setGeometry(100, 100, 1280, 720)
        self._init_ui()
        self.project_manager_panel.projectSelected.connect(self.load_project_images)
        self.project_manager_panel.projectSelected.connect(self.annotation_widget.set_project)
        self._open_default_folders()
        self.connect_offset_sync()
        self.connect_project_manager()

    def _init_ui(self):
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        
        # Create tab widget
        self.tabs = QTabWidget()
        
        # 1. 项目管理面板
        self.project_manager_panel = ProjectManagerPanel()
        self.tabs.addTab(self.project_manager_panel, "项目管理")
        
        # 2. 标注工具
        self.annotation_widget = AnnotationWidget()
        self.tabs.addTab(self.annotation_widget, "标注工具")
        
        # 3. 对齐工具
        self.alignment_widget = AlignmentWidget()
        self.tabs.addTab(self.alignment_widget, "对齐工具")
        
        # 移除预览视图
        # self.preview_view = ImageView("预览视图")
        # self.tabs.addTab(self.preview_view, "预览视图")
        
        # 4. 抽帧工具 (之前是5)
        self.frame_extraction_panel = FrameExtractionPanel()
        self.tabs.addTab(self.frame_extraction_panel, "抽帧工具")
        
        layout = QVBoxLayout()
        layout.addWidget(self.tabs)
        main_widget.setLayout(layout)

    def _open_default_folders(self):
        rgb_folder = r"F:\hzx-uav\runs\V"
        ti_folder = r"F:\hzx-uav\runs\T"
        if rgb_folder and ti_folder:
            self.rgb_image_list = sorted(glob.glob(os.path.join(rgb_folder, "*.[jJ][pP][gG]")) +
                                          glob.glob(os.path.join(rgb_folder, "*.[pP][nN][gG]")))
            self.ti_image_list = sorted(glob.glob(os.path.join(ti_folder, "*.[jJ][pP][gG]")) +
                                         glob.glob(os.path.join(ti_folder, "*.[pP][nN][gG]")))
            if self.rgb_image_list and self.ti_image_list:
                rgb_path = self.rgb_image_list[0]
                ti_path = self.ti_image_list[0]
                self.alignment_widget.load_images(rgb_path, ti_path)
                self.annotation_widget.load_images(rgb_path, ti_path)
                # 移除对预览视图的引用
                # self.preview_view.load_images(rgb_path, ti_path)

    def connect_offset_sync(self):
        self.alignment_widget.offsetChanged.connect(self.syncAlignmentToAnnotator)
        self.annotation_widget.offsetChanged.connect(self.syncAnnotatorToAlignment)

    def syncAlignmentToAnnotator(self, x, y):
        if self.annotation_widget.offset_slider_x.value() != x:
            self.annotation_widget.offset_slider_x.blockSignals(True)
            self.annotation_widget.offset_slider_x.setValue(x)
            self.annotation_widget.offset_slider_x.blockSignals(False)
        if self.annotation_widget.offset_slider_y.value() != y:
            self.annotation_widget.offset_slider_y.blockSignals(True)
            self.annotation_widget.offset_slider_y.setValue(y)
            self.annotation_widget.offset_slider_y.blockSignals(False)
        self.annotation_widget.update_movement()

    def syncAnnotatorToAlignment(self, x, y):
        if self.alignment_widget.offset_slider_x.value() != x:
            self.alignment_widget.offset_slider_x.blockSignals(True)
            self.alignment_widget.offset_slider_x.setValue(x)
            self.alignment_widget.offset_slider_x.blockSignals(False)
        if self.alignment_widget.offset_slider_y.value() != y:
            self.alignment_widget.offset_slider_y.blockSignals(True)
            self.alignment_widget.offset_slider_y.setValue(y)
            self.alignment_widget.offset_slider_y.blockSignals(False)
        self.alignment_widget.offset_x_value_label.setText(str(x))
        self.alignment_widget.offset_y_value_label.setText(str(y))
        self.alignment_widget.update_alignment()

    def load_project_images(self, project_path):
        project_name = os.path.basename(project_path)
        self.setWindowTitle(f"MMReID标注工具 - 当前项目：{project_name}")
        frames_dir = os.path.join(project_path, "Frames")
        
        # 使用os.walk递归查找所有名为"T"或"V"的目录
        candidates = []
        for root, dirs, _ in os.walk(frames_dir):
            for d in dirs:
                if d in ["T", "V"]:
                    candidates.append(os.path.join(root, d))
        
        if not candidates:
            QMessageBox.warning(self, "警告", "在Frames目录下未找到T/V文件夹")
            return
        
        # 按父目录名称分组，父目录代表一组抽帧结果
        groups = {}
        for folder in candidates:
            parent = os.path.basename(os.path.dirname(folder))
            groups.setdefault(parent, {})[os.path.basename(folder)] = folder
        
        # 选择第一组中包含至少一个模态的文件夹作为候选
        chosen = None
        for key, modalities in groups.items():
            if modalities:
                chosen = modalities
                break
        
        if not chosen:
            QMessageBox.warning(self, "警告", "未找到有效的T/V文件夹")
            return
        
        # 优先选择V和T，若缺失则用现有的模态替代
        rgb_folder = chosen.get("V")
        ti_folder = chosen.get("T")
        if not rgb_folder:
            rgb_folder = ti_folder
        if not ti_folder:
            ti_folder = rgb_folder
        
        self.rgb_image_list = sorted(glob.glob(os.path.join(rgb_folder, "*.[jJ][pP][gG]")) +
                                     glob.glob(os.path.join(rgb_folder, "*.[pP][nN][gG]")))
        self.ti_image_list = sorted(glob.glob(os.path.join(ti_folder, "*.[jJ][pP][gG]")) +
                                    glob.glob(os.path.join(ti_folder, "*.[pP][nN][gG]")))
        if self.rgb_image_list and self.ti_image_list:
            rgb_path = self.rgb_image_list[0]
            ti_path = self.ti_image_list[0]
            self.alignment_widget.load_images(rgb_path, ti_path)
            self.annotation_widget.load_images(rgb_path, ti_path)
            # 移除对预览视图的引用
            # self.preview_view.load_images(rgb_path, ti_path)
        else:
            QMessageBox.warning(self, "警告", "未找到有效的RGB或TI图像")

    def connect_project_manager(self):
        # 连接视频进度面板的去标注信号
        if hasattr(self.project_manager_panel, "video_progress_panel") and \
        hasattr(self.project_manager_panel.video_progress_panel, "gotoAnnotation"):
            self.project_manager_panel.video_progress_panel.gotoAnnotation.connect(self.load_video_for_annotation)

    def load_video_for_annotation(self, video_path):
        """加载视频到标注工具并切换到标注视图"""
        # 切换到标注标签页
        self.tabs.setCurrentIndex(1)  # 使用正确的tabs属性名
        
        # 通知标注工具加载视频
        if hasattr(self, "annotation_widget"):
            self.annotation_widget.load_video(video_path)


